package cn.i7baozh.starter.gateway.loadbalancer;

import cn.i7baozh.starter.gateway.constant.GatewayConstants;
import cn.i7baozh.starter.gateway.loadbalancer.component.AvailableGroupNacosConfigComponent;
import cn.i7baozh.starter.gateway.loadbalancer.dto.InstanceInfo;
import com.alibaba.cloud.commons.lang.StringUtils;
import com.alibaba.cloud.nacos.balancer.NacosBalancer;
import com.alibaba.nacos.api.naming.pojo.Instance;
import com.alibaba.nacos.client.naming.utils.Chooser;
import com.alibaba.nacos.client.naming.utils.CollectionUtils;
import com.alibaba.nacos.client.naming.utils.Pair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.DefaultResponse;
import org.springframework.cloud.client.loadbalancer.EmptyResponse;
import org.springframework.cloud.client.loadbalancer.Request;
import org.springframework.cloud.client.loadbalancer.Response;
import org.springframework.cloud.loadbalancer.core.NoopServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import reactor.core.publisher.Mono;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import static com.alibaba.nacos.client.utils.LogUtils.NAMING_LOGGER;

/**
 * 通过远程配置动态拉取不同分组的服务列表
 *
 * @author : baoqi.zhang
 * @date :   2022/5/12 20:14
 */
public class AvailableGroupLoadBalancer implements ReactorServiceInstanceLoadBalancer {

    private static final Logger log = LoggerFactory.getLogger(AvailableGroupLoadBalancer.class);

    private final String serviceId;

    private ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider;

    @Autowired
    private AvailableGroupNacosConfigComponent availableGroupNacosConfigComponent;

    public AvailableGroupLoadBalancer(
            ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider,
            String serviceId) {
        this.serviceId = serviceId;
        this.serviceInstanceListSupplierProvider = serviceInstanceListSupplierProvider;
    }

    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        ServiceInstanceListSupplier supplier = serviceInstanceListSupplierProvider
                .getIfAvailable(NoopServiceInstanceListSupplier::new);
//        DefaultRequestContext context = (DefaultRequestContext)request.getContext();
//        RequestData requestData = (RequestData) context.getClientRequest();
//        HttpHeaders headers = requestData.getHeaders();
//
//        String envHeader = Optional.ofNullable(headers.getFirst(HeaderLoadBalanceConfigComponent.SELECT_GROUP_HEADER)).orElse(HeaderLoadBalanceConfigComponent.SELECT_GROUP_HEADER);
        //获取当前可用的服务组信息
        String availableServiceGroup = availableGroupNacosConfigComponent.getCurrentAvailableServiceGroup();
//        log.info("envHeader:{}", availableServiceGroup);
        return supplier.get().next().map(serviceInstance -> getInstanceResponse(serviceInstance, availableServiceGroup));
    }

    private Response<ServiceInstance> getInstanceResponse(
            List<ServiceInstance> serviceInstances, String availableServiceGroup) {
        if (serviceInstances.isEmpty()) {
            log.warn("No servers available for service: " + this.serviceId);
            return new EmptyResponse();
        }

        try {

            List<ServiceInstance> instancesToChoose = serviceInstances;
            if (StringUtils.isNotBlank(availableServiceGroup)) {
                List<ServiceInstance> sameClusterInstances = serviceInstances.stream()
                        .filter(serviceInstance -> {
                            String cluster = serviceInstance.getMetadata()
                                    .get(GatewayConstants.SELECT_GROUP_HEADER);

                            return StringUtils.equals(cluster, availableServiceGroup);
                        }).collect(Collectors.toList());
                if (!CollectionUtils.isEmpty(sameClusterInstances)) {
                    instancesToChoose = sameClusterInstances;
                }
            } else {
                log.warn(
                        "A cross-cluster call occurs，name = {}, clusterName = {}, instance = {}",
                        serviceId, availableServiceGroup, serviceInstances);
            }

            return new DefaultResponse(getHostByRandomWeight3(instancesToChoose));
        } catch (Exception e) {
            log.warn("NacosLoadBalancer error", e);
            return null;
        }

    }

    /**
     *
     * Spring Cloud LoadBalancer Choose instance by weight.
     * @see NacosBalancer#getHostByRandomWeight3(java.util.List)
     * @param serviceInstances Instance List
     * @return the chosen instance
     */
    public ServiceInstance getHostByRandomWeight3(
            List<ServiceInstance> serviceInstances) {
        Map<Instance, ServiceInstance> instanceMap = new HashMap<>();
        List<Instance> nacosInstance = serviceInstances.stream().map(serviceInstance -> {
            Map<String, String> metadata = serviceInstance.getMetadata();

            // see
            // com.alibaba.cloud.nacos.discovery.NacosServiceDiscovery.hostToServiceInstance()
            Instance instance = new Instance();
            instance.setIp(serviceInstance.getHost());
            instance.setPort(serviceInstance.getPort());

            /**
             * 通过远程调用获取配置
             */
            InstanceInfo instanceInfo = availableGroupNacosConfigComponent
                    .getInstanceInfo(serviceInstance.getHost(), serviceInstance.getPort());
            if (instanceInfo != null) {
                instance.setWeight(instanceInfo.getWeight());
                instance.setHealthy(instanceInfo.getHealthy());
            } else {
                instance.setWeight(Double.parseDouble(metadata.get("nacos.weight")));
                instance.setHealthy(Boolean.parseBoolean(metadata.get("nacos.healthy")));
            }

            instanceMap.put(instance, serviceInstance);
            return instance;
        }).collect(Collectors.toList());

        Instance instance = getHostByRandomWeight(nacosInstance);
        return instanceMap.get(instance);
    }

    /**
     * Return one host from the host list by random-weight.
     *
     * see Balancer#getHostByRandomWeight(java.util.List)
     * @param hosts The list of the host.
     * @return The random-weight result of the host
     */
    protected static Instance getHostByRandomWeight(List<Instance> hosts) {
        NAMING_LOGGER.debug("entry randomWithWeight");
        if (hosts == null || hosts.size() == 0) {
            NAMING_LOGGER.debug("hosts == null || hosts.size() == 0");
            return null;
        }
        NAMING_LOGGER.debug("new Chooser");
        List<Pair<Instance>> hostsWithWeight = new ArrayList<Pair<Instance>>();
        for (Instance host : hosts) {
            if (host.isHealthy()) {
                hostsWithWeight.add(new Pair<Instance>(host, host.getWeight()));
            }
        }
        NAMING_LOGGER.debug("for (Host host : hosts)");
        Chooser<String, Instance> vipChooser = new Chooser<String, Instance>("www.taobao.com");
        vipChooser.refresh(hostsWithWeight);
        NAMING_LOGGER.debug("vipChooser.refresh");
        return vipChooser.randomWithWeight();
    }
}
